import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

class KMeans:
    def __init__(self, n_clusters=3, max_iterations=100, tolerance=1e-4):
        self.n_clusters = n_clusters  
        self.max_iterations = max_iterations 
        self.tolerance = tolerance 
        self.centroids = None 
        self.labels = None 

    def fit(self, X):

        np.random.seed(42)  
        random_indices = np.random.choice(X.shape[0], self.n_clusters, replace=False)
        self.centroids = X[random_indices]

        for _ in range(self.max_iterations):

            distances = np.linalg.norm(X[:, np.newaxis] - self.centroids, axis=2)

            self.labels = np.argmin(distances, axis=1)

            previous_centroids = self.centroids.copy()

            self.centroids = np.array([X[self.labels == k].mean(axis=0) for k in range(self.n_clusters)])

   
            if np.linalg.norm(self.centroids - previous_centroids) < self.tolerance:
                break

    def predict(self, X):
 
        distances = np.linalg.norm(X[:, np.newaxis] - self.centroids, axis=2)
        return np.argmin(distances, axis=1)

if __name__ == '__main__':
 
    iris_data = pd.read_csv('iris.csv', header=None)
    
    X = iris_data.iloc[:, :-1].values

    kmeans = KMeans(n_clusters=3)

    kmeans.fit(X)

    print("质心:")
    print(kmeans.centroids)
    print("每个样本的标签:")
    print(kmeans.labels)

    plt.figure(figsize=(8, 6))
    plt.scatter(X[:, 0], X[:, 1], c=kmeans.labels, cmap='viridis', marker='o', label='数据点')
    plt.scatter(kmeans.centroids[:, 0], kmeans.centroids[:, 1], c='red', marker='x', s=200, label='质心')
    plt.title('K-means Clustering (Iris Dataset)')
    plt.xlabel('特征 1 (萼片长度)')
    plt.ylabel('特征 2 (萼片宽度)')
    plt.legend()
    plt.grid()
    plt.show()